# FMCW雷达多参数估计数据集类
import os
import torch
from torch.utils.data import Dataset
from apps.fmcw.conf.app_config import AppConfig as AF
from apps.fmcw.core.ira_utils import IraUtils as IU

class FmcwMpeDssV3(Dataset):

    def __init__(self, dss_fn:str = './work/datasets/mpe_train_v3.txt'):
        self.name = 'dss.fmcw_mpe_dss.FmcwMpeDss'
        self.Xs, self.ys = [], []
        num = 0
        with open(dss_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                arrs0 = row.split(',')
                self.Xs.append(arrs0[0])
                self.ys.append(arrs0[1])
        print(f'### self.Xs: {type(self.Xs)}; self.ys: {type(self.ys[0])};')

    def __len__(self):
        return len(self.Xs)

    def __getitem__(self, idx):
        X_fn, y_fn = self.Xs[idx], self.ys[idx]
        X = torch.load(X_fn, weights_only=True)
        fn_stem = y_fn[-9:]
        folder = y_fn[:-10]
        y_p_fn = f'{folder}y_p{fn_stem}'
        y_rvt_fn = f'{folder}y_rvt{fn_stem}'
        y_rvt_base_fn = f'{folder}y_rvt_base{fn_stem}'
        y_p = torch.load(y_p_fn, weights_only=True)
        y_rvt = torch.load(y_rvt_fn, weights_only=True)
        y_rvt_base = torch.load(y_rvt_base_fn, weights_only=True)
        X = X.permute(1,0,2)
        return X, y_p, y_rvt, y_rvt_base
    
    @staticmethod
    def generate_dss_desc() -> None:
        dss_fn = './work/datasets/mpe_train.txt'
        raw_fn = '/home/psdz/diskc/yantao/adev/iragent/work/datasets/mpe'
        fns_dict = {}
        with open(dss_fn, 'w', encoding='utf-8') as wfd:
            for root, dirs, files in os.walk(raw_fn):
                for rfn in files:
                    key = rfn[2:-3]
                    if 'X_' in rfn:
                        print(f'样本数据集文件：{rfn}; key: {rfn[2:-4]}; ffn: {root}/{rfn};')
                        fns_dict[key] = f'{root}/{rfn}'
            for root, dirs, files in os.walk(raw_fn):
                for rfn in files:
                    key = rfn[2:-3]
                    if 'y_' in rfn:
                        print(f'写入文件：{root}/{rfn};')
                        X_fn = fns_dict[key]
                        wfd.write(f'{X_fn},{root}/{rfn}\n')

    @staticmethod
    def prepare_rsp_yolo_ds_v3() -> None:
        # 读入数据集文件列表
        dss_fn = './work/datasets/mpe_train.txt'
        Xs, ys = [], []
        y_dict = {}
        with open(dss_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                arrs0 = row.split(',')
                Xs.append(arrs0[0])
                ys.append(arrs0[1])
                y_dict[arrs0[1].split('/')[-1][1:]] = arrs0[1]
        num = 0
        for x_fn in Xs:
            y_fn = y_dict[x_fn.split('/')[-1][1:]]
            y = torch.load(y_fn, weights_only=True)
            # every bin propability of real target 256*256*8：概率模型
            y_p = torch.zeros((1, AF.numADC, AF.numTX*AF.numRX), dtype=torch.float32)
            # r: range; v: velocity; t: theta; f: phi 256*256*8*3：距离、速度、水平到达角
            y_rvt_base = torch.zeros((3, AF.numADC, AF.numTX*AF.numRX), dtype=torch.float32) # 保存FFT基准值
            y_rvt = torch.zeros((3, AF.numADC, AF.numTX*AF.numRX), dtype=torch.float32) # 保存在基准值上的增量
            for idx in range(0, len(y), 4):
                if y[idx] > 0.9:
                    range_idx = IU.get_near_idx(AF.range_bin, y[idx+1])
                    velocity_idx = IU.get_near_idx(AF.velocity_bin, y[idx+2])
                    theta_idx = IU.get_near_idx(AF.theta_bin, y[idx+3])
                    y_p[0][range_idx][theta_idx] = 1.0
                    # 存储基准值
                    y_rvt_base[0][range_idx][theta_idx] = AF.range_bin[range_idx]
                    y_rvt_base[1][range_idx][theta_idx] = AF.velocity_bin[velocity_idx]
                    y_rvt_base[2][range_idx][theta_idx] = AF.theta_bin[theta_idx]
                    # 存储增量值
                    y_rvt[0][range_idx][theta_idx] = y[idx+1] - AF.range_bin[range_idx]
                    y_rvt[1][range_idx][theta_idx] = y[idx+2] - AF.velocity_bin[velocity_idx]
                    y_rvt[2][range_idx][theta_idx] = y[idx+3] - AF.theta_bin[theta_idx]
                    # y[idx+1:idx+4] = (y[idx+1:idx+4]-AF.Y_MIN[idx+1:idx+4])/(AF.Y_MAX[idx+1:idx+4] - AF.Y_MIN[idx+1:idx+4])
                    # y_rvt[0][range_idx][theta_idx] = y[idx+1]
                    # y_rvt[1][range_idx][theta_idx] = y[idx+2]
                    # y_rvt[2][range_idx][theta_idx] = y[idx+3]
            y_p_fn = f'./work/datasets/mpe_v3/y_p{x_fn.split("/")[-1][1:]}'
            torch.save(y_p, y_p_fn)
            y_rvt_fn = f'./work/datasets/mpe_v3/y_rvt{x_fn.split("/")[-1][1:]}'
            torch.save(y_rvt, y_rvt_fn)
            y_rvt_base_fn = f'./work/datasets/mpe_v3/y_rvt_base{x_fn.split("/")[-1][1:]}'
            torch.save(y_rvt_base, y_rvt_base_fn)
            num += 1
            print(f'进度：{x_fn.split("/")[-1][2:]} {num} / {len(Xs)}; ')

        
